# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: yuxuan

import numpy as np

from coin.base.book.types import TradeSide


def sampler_map():
  return {
      'SweepSampler': SweepSampler,
      'SweepInferredFromTradeSampler': SweepInferredFromTradeSampler,
      'L1PriceMovedSampler': L1PriceMovedSampler,
      'TradeExceedBookSampler': TradeExceedBookSampler
  }


def get_sampler_from_config(config, parent, products):
  if config:
    if not isinstance(config, list):
      config = [config]
    sampler = []
    for cfg in config:
      cls = sampler_map()[cfg['name']]
      idx = cfg.get('exchanges', [0, -1])
      product_param = [products[i] for i in idx]
      params = [parent] + product_param + cfg['params']
      sampler.append(cls(*params))
    if len(sampler) == 1:
      return sampler[0]
    else:
      return MultipleSampler(sampler)
  else:
    return None


class Sampler(object):
  def __init__(self, parent, product0, product1):
    self._parent = parent
    self._product0 = product0
    self._product1 = product1

  def get_ts(self):
    return self._parent.get_ts()

  def get_feed(self, product):
    return self._parent._last_feed[product]

  def on_book(self, product, book):
    pass

  def on_trade(self, product, trade):
    pass

  def sample(self):
    raise NotImplementedError()


class MultipleSampler(Sampler):
  def __init__(self, samplers):
    assert len(samplers) > 1, samplers
    self._samplers = samplers

  def on_book(self, product, book):
    for sampler in self._samplers:
      sampler.on_book(product, book)

  def on_trade(self, product, trade):
    for sampler in self._samplers:
      sampler.on_trade(product, trade)

  def sample(self):
    samples = [sampler.sample() for sampler in self._samplers]
    return np.any(samples)


class TradeExceedBookSampler(Sampler):
  def __init__(self, parent, product0, product1, threshold_bp):
    Sampler.__init__(self, parent, product0, product1)
    assert threshold_bp > 0
    self._threshold = threshold_bp * 1e-4
    self._last_buy_sampled_step = 0
    self._last_sell_sampled_step = 0
    self._last_bid0 = 0
    self._last_ask0 = 0
    self._sample = False

  def on_book(self, product, book):
    if product == self._product0:
      bids, asks = self.get_feed(self._product0)
      bid0 = bids[0][0]
      ask0 = asks[0][0]
      if not np.isclose(ask0, self._last_ask0):
        self._last_buy_sampled_step = 0
      if not np.isclose(bid0, self._last_bid0):
        self._last_sell_sampled_step = 0
      self._last_bid0 = bid0
      self._last_ask0 = ask0

  def on_trade(self, product, trade):
    if product == self._product0:
      if trade.side == TradeSide.TRADE_BUY_SIDE:
        last_ask0 = self._last_ask0 or trade.price
        level = int((trade.price / last_ask0 - 1) / self._threshold)
        if level > self._last_buy_sampled_step:
          self._last_buy_sampled_step = level
          self._sample = True
      else:
        last_bid0 = self._last_bid0 or trade.price
        level = int((last_bid0 / trade.price - 1) / self._threshold)
        if level > self._last_sell_sampled_step:
          self._last_sell_sampled_step = level
          self._sample = True

  def sample(self):
    if self._sample:
      self._sample = False
      return True
    return False


class SweepSampler(Sampler):
  def __init__(self, parent, product0, product1, sweep_levels):
    Sampler.__init__(self, parent, product0, product1)
    self._sweep_levels = sweep_levels
    self._bids, self._asks = None, None
    self._sample = False

  def on_book(self, product, book):
    if product == self._product0:
      bids, asks = self.get_feed(product)
      if self._bids is not None:
        if (bids[0][0] < self._bids[self._sweep_levels - 1][0] - 1e-6) \
              or (asks[0][0] > self._asks[self._sweep_levels - 1][0] + 1e-6):
          self._sample = True
        else:
          self._sample = False
      self._bids, self._asks = bids, asks

  def sample(self):
    if self._sample:
      self._sample = False
      return True
    return False


class SweepInferredFromTradeSampler(Sampler):
  def __init__(self, parent, product0, product1, sweep_levels):
    Sampler.__init__(self, parent, product0, product1)
    self._sweep_levels = sweep_levels
    self._bids, self._asks = None, None
    self._book_ts = 0
    self._sample = False
    self._last_sample_book_ts = -1

  def on_book(self, product, book):
    if product == self._product0:
      self._bids, self._asks = self.get_feed(product)
      self._book_ts = self.get_ts()

  def on_trade(self, product, trade):
    if product == self._product0:
      sample = False
      if self._bids is not None:
        if trade.price > self._asks[self._sweep_levels - 1][0] \
              or trade.price < self._bids[self._sweep_levels - 1][0]:
          if self._last_sample_book_ts != self._book_ts:
            sample = True
            self._last_sample_book_ts = self._book_ts
      self._sample = sample

  def sample(self):
    if self._sample:
      self._sample = False
      return True
    return False


class L1PriceMovedSampler(Sampler):
  def __init__(self, parent, product0, product1, threshold_bp):
    Sampler.__init__(self, parent, product0, product1)
    self._threshold = threshold_bp * 1e-4
    self._bids, self._asks = None, None
    self._sample = False

  def on_book(self, product, book):
    if product == self._product0:
      sample = False
      if self._bids is not None:
        bids, asks = self.get_feed(product)
        threshold = (bids[0][0] + asks[0][0]) / 2 * self._threshold
        if abs(bids[0][0] - self._bids[0][0]) > threshold \
                or abs(asks[0][0] - self._asks[0][0]) > threshold:
          sample = True
      self._bids, self._asks = self.get_feed(product)
      self._sample = sample

  def sample(self):
    if self._sample:
      self._sample = False
      return True
    return False
